{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nova Data Mixing\n",
    "\n",
    "Data mixing blends your custom training data with Nova's curated synthetic datasets\n",
    "(code, math, chat, planning, instruction-following, reasoning, etc.) to prevent\n",
    "catastrophic forgetting while specializing the model on your domain.\n",
    "\n",
    "> **Important:** Data mixing is only supported with **serverless** compute type.\n",
    "> It is not available for serverful training jobs (SMTJ) or HyperPod clusters.\n",
    "\n",
    "## What you will learn\n",
    "\n",
    "1. Configure `DataMixingConfig` with customer and Nova data percentages\n",
    "2. Create an `SFTTrainer` with data mixing enabled\n",
    "3. Set hyperparameters and submit a training job\n",
    "4. Monitor job status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import json\n",
    "import boto3"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# === Fill in your AWS resources ===\n",
    "REGION = \"<your-region>\"  # e.g. \"us-east-1\"\n",
    "ROLE_ARN = \"<your-execution-role-arn>\"\n",
    "S3_BUCKET = \"<your-s3-bucket>\"  # e.g. \"sagemaker-us-east-1-123456789012\"\n",
    "\n",
    "S3_OUTPUT_PATH = f\"s3://{S3_BUCKET}/sft-data-mixing/output\"\n",
    "TRAINING_DATASET = f\"s3://{S3_BUCKET}/datasets/sft_training_data.jsonl\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Configure Data Mixing\n",
    "\n",
    "Data mixing controls the blend between your custom training data and Nova's internal\n",
    "curated datasets. `customer_data_percent` sets how much of the training data comes from\n",
    "your dataset. The remaining portion is distributed among Nova categories according to\n",
    "`nova_data_percentages`.\n",
    "\n",
    "Available Nova categories include: `code`, `math`, `chat`, `planning`,\n",
    "`instruction-following`, `reasoning`, `stem`, `rag`, `factuality`, etc."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.train.data_mixing_config import DataMixingConfig\n",
    "\n",
    "# 70% of training data from your dataset, 30% from Nova curated data\n",
    "# Within Nova data: 30% code, 70% math\n",
    "data_mixing_config = DataMixingConfig(\n",
    "    customer_data_percent=70.0,\n",
    "    nova_data_percentages={\n",
    "        \"code\": 30.0,\n",
    "        \"math\": 70.0,\n",
    "    },\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create SFTTrainer with Data Mixing\n",
    "\n",
    "Pass the `DataMixingConfig` to `SFTTrainer`. Since data mixing only works with\n",
    "serverless compute, no `compute` parameter is needed."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.train import SFTTrainer\n",
    "from sagemaker.train.common import TrainingType\n",
    "\n",
    "sft_trainer = SFTTrainer(\n",
    "    model=\"amazon.nova-2-lite-v1\",\n",
    "    training_type=TrainingType.LORA,\n",
    "    training_dataset=TRAINING_DATASET,\n",
    "    s3_output_path=S3_OUTPUT_PATH,\n",
    "    role=ROLE_ARN,\n",
    "    data_mixing_config=data_mixing_config,\n",
    "    base_job_name=\"sft-datamix\",\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Set Hyperparameters and Submit Training Job"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Set hyperparameters\n",
    "sft_trainer.hyperparameters.max_steps = 50\n",
    "sft_trainer.hyperparameters.learning_rate = 5e-6\n",
    "sft_trainer.hyperparameters.global_batch_size = 32\n",
    "\n",
    "# Submit (non-blocking)\n",
    "training_job = sft_trainer.train(wait=False)\n",
    "print(f\"Training job submitted: {training_job}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Monitor Training Job"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from sagemaker.core.resources import TrainingJob\n",
    "\n",
    "job = TrainingJob.get(training_job_name=training_job.training_job_name)\n",
    "print(f\"Status: {job.training_job_status}\")\n",
    "print(f\"Secondary Status: {job.secondary_status}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Alternative: Different Data Mix Configurations\n",
    "\n",
    "Here are some common configuration patterns depending on your use case."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# High specialization: mostly your data\n",
    "high_specialization = DataMixingConfig(\n",
    "    customer_data_percent=90.0,\n",
    "    nova_data_percentages={\n",
    "        \"reasoning\": 100.0,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Balanced: equal split with multiple Nova categories\n",
    "balanced_mix = DataMixingConfig(\n",
    "    customer_data_percent=50.0,\n",
    "    nova_data_percentages={\n",
    "        \"code\": 40.0,\n",
    "        \"reasoning\": 30.0,\n",
    "        \"math\": 30.0,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Light specialization: preserve broad capabilities\n",
    "light_specialization = DataMixingConfig(\n",
    "    customer_data_percent=30.0,\n",
    "    nova_data_percentages={\n",
    "        \"code\": 25.0,\n",
    "        \"math\": 25.0,\n",
    "        \"chat\": 25.0,\n",
    "        \"reasoning\": 25.0,\n",
    "    },\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tips\n",
    "\n",
    "- **High customer_data_percent (80\u201390%)** \u2014 Use when your task is well-defined and you have enough data.\n",
    "- **Balanced (50\u201370%)** \u2014 Good default for most use cases.\n",
    "- **Low customer_data_percent (20\u201340%)** \u2014 Preserve base model capabilities with light specialization.\n",
    "- **Nova category selection** \u2014 Choose categories that complement your task (e.g., `code` + `reasoning` for a coding assistant)."
   ]
  }
 ]
}